import cv2
import torch
import pickle
import numpy as np
import matplotlib.pyplot as plt
import trimesh
from torchvision import transforms
import os

from plyfile import PlyData


# mask = cv2.imread('/media/lab502/18T/ccc/clean_pvnet/clean-pvnet/data/custom_sunhan/mask/000000.png')
# mask = cv2.imread('/media/lab502/18T/ccc/clean_pvnet/clean-pvnet/data/linemod/custom_cat/mask/0.png')
#
# mask[mask>0] = 255
# cv2.imshow('mask',mask)
# cv2.waitKey()
# pass




# Pickle functions to save and load dictionaries
def save_obj(obj, name):
    with open(name + '.pkl', 'wb') as f:
        pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL)


def load_obj(name):
    with open(name + '.pkl', 'rb') as f:
        return pickle.load(f)

# helper function to plot grpahs


def visualize(array):
    "Plot all images in the array of tensors in one row"
    for z in range(0, len(array)):
        temp = array[z]
        if temp.ndim > 3:  # tensor output in the form NCHW
            temp = (torch.argmax(temp, dim=1).squeeze())
        if len(temp.shape) >= 3:
            plt.figure()
            plt.imshow(np.transpose(
                temp.detach().numpy().squeeze(), (1, 2, 0)))
            plt.show()
        else:
            plt.figure()
            plt.imshow(temp.detach().numpy(), cmap='gray')


model_path = os.path.join('/media/lab502/35B489DB53B3E4C2/datasets/custom/model.ply')
ply = PlyData.read(model_path)
data = ply.elements[0].data
x = data['x']
y = data['y']
z = data['z']
print(np.min(x), np,min(y), np,min(z),  np.max(x), np,max(y), np,max(z))
def compute_projection(points_3D,internal_calibration):
    points_3D = points_3D.T
    projections_2d = np.zeros((2, points_3D.shape[1]), dtype='float32')
    camera_projection = (internal_calibration).dot(points_3D)
    projections_2d[0, :] = camera_projection[0, :]/camera_projection[2, :]
    projections_2d[1, :] = camera_projection[1, :]/camera_projection[2, :]
    return projections_2d
def create_bounding_box(img, data,pose, intrinsic_matrix,color=(255,255,255)):
    "Create a bounding box around the object"
    # 8 corner points of the ptcld data

    #
    # min_x, min_y, min_z =  -79.27367, -73.72931,269.4129
    # max_x, max_y, max_z =  103.2678 , -14.11505, 366.5242
    x = data['x']
    y = data['y']
    z = data['z']
    x_min = np.min(x)
    y_min = np.min(y)
    z_min = np.min(z)
    x_max = np.max(x)
    y_max = np.max(y)
    z_max = np.max(z)
    '''
    min_x, min_y, min_z =  -79.27367, -73.72931,269.4129
    max_x, max_y, max_z =  103.2678 , -14.11505, 366.5242
    '''
    min_x, min_y, min_z = x_min, y_min, z_min
    max_x, max_y, max_z = x_max, y_max, z_max
    corners_3D = np.array([[max_x, min_y, min_z],
                           [max_x, min_y, max_z],
                           [min_x, min_y, max_z],
                           [min_x, min_y, min_z],
                           [max_x, max_y, min_z],
                           [max_x, max_y, max_z],
                           [min_x, max_y, max_z],
                           [min_x, max_y, min_z]])
    # print(corners_3D)

    # convert these 8 3D corners to 2D points
    ones = np.ones((corners_3D.shape[0], 1))
    homogenous_coordinate = np.append(corners_3D, ones, axis=1)

    # Perspective Projection to obtain 2D coordinates for masks
    homogenous_2D = intrinsic_matrix @ (pose @ homogenous_coordinate.T)
    coord_2D = homogenous_2D[:2, :] / homogenous_2D[2, :]
    coord_2D = ((np.floor(coord_2D)).T).astype(int)
    # coord_2D = ((np.floor(coord_2D))).astype(int)
    # Draw lines between these 8 points

    '''
    # one = np.array([0,0,0,1])
    # pose1 = np.vstack((pose,one))
    corners_3D_trans = trimesh.transformations.transform_points(corners_3D,np.linalg.inv(pose))
    coord_2D =  compute_projection(corners_3D_trans,intrinsic_matrix)
    coord_2D = coord_2D.T
    coord_2D[:, 0] = coord_2D[:, 0]
    coord_2D[:, 1] = coord_2D[:, 1]
    coord_2D = ((np.floor(coord_2D))).astype(int)
    '''
    img = cv2.line(img, tuple(coord_2D[0]), tuple(coord_2D[1]), color, 2)
    img = cv2.line(img, tuple(coord_2D[0]), tuple(coord_2D[3]), color, 2)
    img = cv2.line(img, tuple(coord_2D[0]), tuple(coord_2D[4]), color, 2)
    img = cv2.line(img, tuple(coord_2D[1]), tuple(coord_2D[2]), color, 2)
    img = cv2.line(img, tuple(coord_2D[1]), tuple(coord_2D[5]), color, 2)
    img = cv2.line(img, tuple(coord_2D[2]), tuple(coord_2D[3]), color, 2)
    img = cv2.line(img, tuple(coord_2D[2]), tuple(coord_2D[6]), color, 2)
    img = cv2.line(img, tuple(coord_2D[3]), tuple(coord_2D[7]), color, 2)
    img = cv2.line(img, tuple(coord_2D[4]), tuple(coord_2D[7]), color, 2)
    img = cv2.line(img, tuple(coord_2D[4]), tuple(coord_2D[5]), color, 2)
    img = cv2.line(img, tuple(coord_2D[5]), tuple(coord_2D[6]), color, 2)
    img = cv2.line(img, tuple(coord_2D[6]), tuple(coord_2D[7]), color, 2)
    return img

def create_coor(img, pose, intrinsic_matrix):
    "Create a bounding box around the object"
    # 8 corner points of the ptcld data

    x_color = (255, 0, 0)
    y_color = (0, 255, 0)
    z_color = (0, 0, 255)
    corners_3D = np.array([[0, 0, 0],
                           [100, 0, 0],
                           [0, 100, 0],
                           [0, 0, 100]])

    # convert these 8 3D corners to 2D points
    ones = np.ones((corners_3D.shape[0], 1))
    homogenous_coordinate = np.append(corners_3D, ones, axis=1)

    # Perspective Projection to obtain 2D coordinates for masks
    homogenous_2D = intrinsic_matrix @ (pose @ homogenous_coordinate.T)
    coord_2D = homogenous_2D[:2, :] / homogenous_2D[2, :]
    coord_2D = ((np.floor(coord_2D)).T).astype(int)
    # coord_2D = ((np.floor(coord_2D))).astype(int)
    # Draw lines between these 8 points
    img = cv2.line(img, tuple(coord_2D[0]), tuple(coord_2D[1]), x_color, 5)
    img = cv2.line(img, tuple(coord_2D[0]), tuple(coord_2D[2]), y_color, 5)
    img = cv2.line(img, tuple(coord_2D[0]), tuple(coord_2D[3]), z_color, 5)

    return img

def ADD_score(pt_cld, true_pose, pred_pose, diameter):
    "Evaluation metric - ADD score"
    pred_pose[0:3, 0:3][np.isnan(pred_pose[0:3, 0:3])] = 1
    pred_pose[:, 3][np.isnan(pred_pose[:, 3])] = 0
    target = pt_cld @ true_pose[0:3, 0:3] + np.array(
        [true_pose[0, 3], true_pose[1, 3], true_pose[2, 3]])
    output = pt_cld @ pred_pose[0:3, 0:3] + np.array(
        [pred_pose[0, 3], pred_pose[1, 3], pred_pose[2, 3]])
    avg_distance = (np.linalg.norm(output - target))/pt_cld.shape[0]
    threshold = diameter * 0.1
    # if avg_distance <= threshold:
    #     return avg_distance
    # else:
    #     return 0
    return avg_distance


for i in range(1000):
    print("i:",i)
    image = cv2.imread('/media/lab502/35B489DB53B3E4C2/datasets/custom/rgb/{}.jpg'.format(i))
    mask = cv2.imread('/media/lab502/35B489DB53B3E4C2/datasets/custom/mask/{}.png'.format(i))

    image[mask == 0] = 0

    pose = np.load('/media/lab502/35B489DB53B3E4C2/datasets/custom/pose/pose{}.npy'.format(i))
    # pose = pose.T
    pose2 = pose[0:3, :]
    # print("pose2: ",pose2)
    # pose2[:,3] = pose2[:,3] #* 1000
    # print("pose: ",pose)
    K = np.array([[6.067540283203125000e+02, 0.000000000000000000e+00, 3.231988525390625000e+02],
                  [0.000000000000000000e+00, 6.068151245117188000e+02, 2.386437530517578000e+02],
                  [0.000000000000000000e+00, 0.000000000000000000e+00, 1.000000000000000000e+00]])
    '''
    K = np.array([[6.084540283203125000e+02, 0.000000000000000000e+00, 3.181988525390625000e+02],
                  [0.000000000000000000e+00,6.081151245117188000e+02,2.426437530517578000e+02],
                  [0.000000000000000000e+00,0.000000000000000000e+00,1.000000000000000000e+00]])
    '''

    img = create_bounding_box(image, data, pose2, K)
    cv2.imshow('img', img)
    cv2.waitKey()

